{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "``brainpy.state`` represents a complete architectural redesign built on top of the ``brainstate`` framework. This document explains the design principles and architectural components that make ``brainpy.state`` powerful and flexible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Design Philosophy\n",
    "\n",
    "``brainpy.state``  is built around several core principles:\n",
    "\n",
    "**State-Based Programming**\n",
    "   All dynamical variables are managed as explicit states, enabling automatic differentiation, efficient compilation, and clear data flow.\n",
    "\n",
    "**Modular Composition**\n",
    "   Complex models are built by composing simple, reusable components. Each component has a well-defined interface and responsibility.\n",
    "\n",
    "**Scientific Accuracy**\n",
    "   Integration with ``brainunit`` ensures physical correctness and prevents unit-related errors.\n",
    "\n",
    "**Performance by Default**\n",
    "   JIT compilation and optimization are built into the framework, not an afterthought.\n",
    "\n",
    "**Extensibility**\n",
    "   Adding new neuron models, synapse types, or learning rules is straightforward and follows clear patterns."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Architectural Layers\n",
    "\n",
    "brainpy.state  is organized into several layers:\n",
    "\n",
    "```text\n",
    "┌─────────────────────────────────────────┐\n",
    "│         User Models & Networks          │  ← Your code\n",
    "├─────────────────────────────────────────┤\n",
    "│      BrainPy Components Layer           │  ← Neurons, Synapses, Projections\n",
    "├─────────────────────────────────────────┤\n",
    "│       BrainState Framework              │  ← State management, compilation\n",
    "├─────────────────────────────────────────┤\n",
    "│       JAX + XLA Backend                 │  ← JIT compilation, autodiff\n",
    "└─────────────────────────────────────────┘\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. JAX + XLA Backend\n",
    "\n",
    "The foundation layer provides:\n",
    "\n",
    "- Just-In-Time (JIT) compilation\n",
    "- Automatic differentiation\n",
    "- Hardware acceleration (CPU/GPU/TPU)\n",
    "- Functional transformations (vmap, grad, etc.)\n",
    "\n",
    "### 2. BrainState Framework\n",
    "\n",
    "Built on JAX, ``brainstate`` provides:\n",
    "\n",
    "- State management system\n",
    "- Module composition\n",
    "- Compilation and optimization\n",
    "- Program transformations (for_loop, etc.)\n",
    "\n",
    "### 3. BrainPy Components\n",
    "\n",
    "High-level neuroscience-specific components:\n",
    "\n",
    "- Neuron models (LIF, ALIF, etc.)\n",
    "- Synapse models (Expon, Alpha, etc.)\n",
    "- Projection architectures\n",
    "- Learning rules and plasticity\n",
    "\n",
    "### 4. User Models\n",
    "\n",
    "Your custom networks and experiments built using BrainPy components."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## State Management System\n",
    "\n",
    "### The Foundation: ``brainstate.State``\n",
    "\n",
    "Everything in ``brainpy.state``  revolves around **states**:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.201528Z",
     "start_time": "2025-11-13T09:31:02.936126Z"
    }
   },
   "source": [
    "import brainpy\n",
    "import brainstate\n",
    "import braintools\n",
    "import brainunit as u\n",
    "import jax.numpy as jnp\n",
    "\n",
    "# Create a state\n",
    "voltage = brainstate.State(0.0)  # Single value\n",
    "weights = brainstate.State([[0.1, 0.2], [0.3, 0.4]])  # Matrix"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "States are special containers that:\n",
    "\n",
    "- Track their values across time\n",
    "- Support automatic differentiation\n",
    "- Enable efficient compilation\n",
    "- Handle batching automatically"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### State Types\n",
    "\n",
    "BrainPy uses different state types for different purposes:\n",
    "\n",
    "**ParamState** - Trainable Parameters\n",
    "   Used for weights, time constants, and other trainable parameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.232382Z",
     "start_time": "2025-11-13T09:31:08.227046Z"
    }
   },
   "source": [
    "class MyNeuron(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.tau = brainstate.ParamState(10.0)  # Trainable\n",
    "        self.weight = brainstate.ParamState([[0.1, 0.2]])"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**ShortTermState** - Temporary Variables\n",
    "   Used for membrane potentials, synaptic currents, and other dynamics."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.256361Z",
     "start_time": "2025-11-13T09:31:08.248156Z"
    }
   },
   "source": [
    "class MyNeuron(brainstate.nn.Module):\n",
    "    def __init__(self, size):\n",
    "        super().__init__()\n",
    "        self.V = brainstate.ShortTermState(jnp.zeros(size))  # Dynamic\n",
    "        self.spike = brainstate.ShortTermState(jnp.zeros(size))"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### State Initialization\n",
    "\n",
    "States can be initialized with various strategies:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.332155Z",
     "start_time": "2025-11-13T09:31:08.288203Z"
    }
   },
   "source": [
    "# Define example size and shape\n",
    "size = 100  # Number of neurons\n",
    "shape = (100, 50)  # Weight matrix shape\n",
    "\n",
    "# Constant initialization\n",
    "V = brainstate.ShortTermState(\n",
    "    braintools.init.Constant(-65.0, unit=u.mV)(size)\n",
    ")\n",
    "\n",
    "# Normal distribution\n",
    "V = brainstate.ShortTermState(\n",
    "    braintools.init.Normal(-65.0, 5.0, unit=u.mV)(size)\n",
    ")\n",
    "\n",
    "# Uniform distribution\n",
    "weights = brainstate.ParamState(\n",
    "    braintools.init.Uniform(0.0, 1.0)(shape)\n",
    ")"
   ],
   "outputs": [],
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module System\n",
    "\n",
    "### Base Class: brainstate.nn.Module\n",
    "\n",
    "All BrainPy components inherit from ``brainstate.nn.Module``:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.344089Z",
     "start_time": "2025-11-13T09:31:08.338163Z"
    }
   },
   "source": [
    "class MyComponent(brainstate.nn.Module):\n",
    "    def __init__(self, size):\n",
    "        super().__init__()\n",
    "        # Initialize states\n",
    "        self.state1 = brainstate.ShortTermState(jnp.zeros(size))\n",
    "        self.param1 = brainstate.ParamState(jnp.ones(size))\n",
    "\n",
    "    def update(self, input):\n",
    "        # Define dynamics\n",
    "        pass"
   ],
   "outputs": [],
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Benefits of Module:\n",
    "\n",
    "- Automatic state registration\n",
    "- Nested module support\n",
    "- State collection and filtering\n",
    "- Serialization support"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Module Composition\n",
    "\n",
    "Modules can contain other modules:\n",
    "\n",
    "```python\n",
    "\n",
    "class Network(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.neurons = brainpy.state.LIF(100, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "        self.synapse = brainpy.state.Expon(100, tau=5*u.ms)\n",
    "        self.projection = brainpy.state.AlignPostProj(...)  # Example - requires more setup\n",
    "\n",
    "    def update(self, input):\n",
    "        # Compose behavior\n",
    "        self.projection(spikes)  # Example\n",
    "        self.neurons(input)\n",
    "\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Component Architecture\n",
    "\n",
    "### Neurons\n",
    "\n",
    "Neurons model the dynamics of neural populations:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.354960Z",
     "start_time": "2025-11-13T09:31:08.350705Z"
    }
   },
   "source": [
    "class Neuron(brainstate.nn.Module):\n",
    "    def __init__(self, size, **kwargs):\n",
    "        super().__init__()\n",
    "        # Membrane potential\n",
    "        self.V = brainstate.ShortTermState(jnp.zeros(size))\n",
    "        # Spike output\n",
    "        self.spike = brainstate.ShortTermState(jnp.zeros(size))\n",
    "\n",
    "    def update(self, input_current):\n",
    "        # Update membrane potential\n",
    "        # Generate spikes\n",
    "        pass"
   ],
   "outputs": [],
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Key responsibilities:\n",
    "\n",
    "- Maintain membrane potential\n",
    "- Generate spikes when threshold is crossed\n",
    "- Reset after spiking\n",
    "- Integrate input currents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Synapses\n",
    "\n",
    "Synapses model temporal filtering of spike trains:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.365530Z",
     "start_time": "2025-11-13T09:31:08.361447Z"
    }
   },
   "source": [
    "class Synapse(brainstate.nn.Module):\n",
    "    def __init__(self, size, tau, **kwargs):\n",
    "        super().__init__()\n",
    "        # Synaptic conductance/current\n",
    "        self.g = brainstate.ShortTermState(jnp.zeros(size))\n",
    "        self.tau = tau\n",
    "\n",
    "    def update(self, spike_input):\n",
    "        # Update synaptic variable\n",
    "        # Return filtered output\n",
    "        pass"
   ],
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Key responsibilities:\n",
    "\n",
    "- Filter spike inputs temporally\n",
    "- Model synaptic dynamics (exponential, alpha, etc.)\n",
    "- Provide smooth currents to postsynaptic neurons"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Projections: The Comm-Syn-Out Pattern\n",
    "\n",
    "Projections connect populations using a three-stage architecture:\n",
    "\n",
    "```text\n",
    "Presynaptic Spikes → [Comm] → [Syn] → [Out] → Postsynaptic Neurons\n",
    "                      │         │       │\n",
    "                  Connectivity  │    Current\n",
    "                  & Weights   Dynamics  Injection\n",
    "```\n",
    "\n",
    "**Communication (Comm)**\n",
    "   Handles spike transmission, connectivity, and weights."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.380053Z",
     "start_time": "2025-11-13T09:31:08.371055Z"
    }
   },
   "source": [
    "# Define population sizes\n",
    "pre_size = 100\n",
    "post_size = 50\n",
    "\n",
    "# Define prob and weight\n",
    "prob=0.1\n",
    "weight=0.5\n",
    "\n",
    "comm = brainstate.nn.EventFixedProb(\n",
    "    pre_size, post_size, prob, weight\n",
    ")"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Synaptic Dynamics (Syn)**\n",
    "   Temporal filtering of transmitted spikes."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.390544Z",
     "start_time": "2025-11-13T09:31:08.386339Z"
    }
   },
   "source": [
    "post_size = 50  # Postsynaptic population size\n",
    "\n",
    "syn = brainpy.state.Expon(post_size, tau=5*u.ms)"
   ],
   "outputs": [],
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Output Mechanism (Out)**\n",
    "   How synaptic variables affect postsynaptic neurons."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.400227Z",
     "start_time": "2025-11-13T09:31:08.397246Z"
    }
   },
   "source": [
    "# Current-based output\n",
    "out = brainpy.state.CUBA()  \n",
    "\n",
    "# Or conductance-based output\n",
    "out = brainpy.state.COBA(E=0*u.mV)"
   ],
   "outputs": [],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Complete Projection**"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.420603Z",
     "start_time": "2025-11-13T09:31:08.405533Z"
    }
   },
   "source": [
    "# Define postsynaptic neurons\n",
    "postsynaptic_neurons = brainpy.state.LIF(50, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "\n",
    "# Create complete projection\n",
    "projection = brainpy.state.AlignPostProj(\n",
    "    comm=comm,\n",
    "    syn=syn,\n",
    "    out=out,\n",
    "    post=postsynaptic_neurons\n",
    ")"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This separation provides:\n",
    "\n",
    "- Clear responsibility boundaries\n",
    "- Easy component swapping\n",
    "- Reusable building blocks\n",
    "- Better testing and debugging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compilation and Execution\n",
    "\n",
    "### Time-Stepped Simulation\n",
    "\n",
    "BrainPy uses discrete time steps:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:08.937334Z",
     "start_time": "2025-11-13T09:31:08.430048Z"
    }
   },
   "source": [
    "# Example: create a simple network\n",
    "class SimpleNetwork(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.neurons = brainpy.state.LIF(100, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "    \n",
    "    def update(self, t, i):\n",
    "        # Generate constant input current\n",
    "        inp = jnp.ones(100) * 5.0 * u.nA\n",
    "        with brainstate.environ.context(t=t, i=i):\n",
    "            self.neurons(inp)\n",
    "            return self.neurons.get_spike()\n",
    "\n",
    "network = SimpleNetwork()\n",
    "brainstate.nn.init_all_states(network)\n",
    "\n",
    "# Set global time step\n",
    "brainstate.environ.set(dt=0.1 * u.ms)\n",
    "\n",
    "# Define simulation duration\n",
    "times = u.math.arange(0*u.ms, 1000*u.ms, brainstate.environ.get_dt())\n",
    "indices = u.math.arange(times.size)\n",
    "\n",
    "# Run simulation\n",
    "results = brainstate.transform.for_loop(\n",
    "    network.update,\n",
    "    times,\n",
    "    indices,\n",
    "    pbar=brainstate.transform.ProgressBar(10)\n",
    ")"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "  0%|          | 0/10000 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "53f4cd22f9f54b9494d7457134633d41"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### JIT Compilation\n",
    "\n",
    "Functions are compiled for performance:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.104741Z",
     "start_time": "2025-11-13T09:31:08.980341Z"
    }
   },
   "source": [
    "# Create example input\n",
    "input_example = jnp.ones(100) * 2.0 * u.nA\n",
    "\n",
    "@brainstate.transform.jit\n",
    "def simulate_step(t, i, input_current):\n",
    "    with brainstate.environ.context(t=t, i=i):\n",
    "        return network.update(t, i)\n",
    "\n",
    "# First call: compile\n",
    "result = simulate_step(0.0*u.ms, 0, input_example)  # Slow (compilation)\n",
    "\n",
    "# Subsequent calls: fast\n",
    "result = simulate_step(0.1*u.ms, 1, input_example)  # Fast (compiled)"
   ],
   "outputs": [],
   "execution_count": 13
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compilation benefits:\n",
    "\n",
    "- 10-100x speedup over Python\n",
    "- Automatic GPU/TPU dispatch\n",
    "- Memory optimization\n",
    "- Fusion of operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradient Computation\n",
    "\n",
    "For training, gradients are computed automatically:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.356118Z",
     "start_time": "2025-11-13T09:31:09.113521Z"
    }
   },
   "source": [
    "# Example: Define mock functions for demonstration\n",
    "def compute_loss(predictions, targets):\n",
    "    return jnp.mean((predictions.astype(float) - targets) ** 2)\n",
    "\n",
    "# Mock targets\n",
    "num_steps = 100\n",
    "targets = jnp.zeros((num_steps, 100))\n",
    "\n",
    "def loss_fn():\n",
    "    # Run network for multiple timesteps\n",
    "    def step(t, i):\n",
    "        with brainstate.environ.context(t=t, i=i):\n",
    "            return network.update(t, i)\n",
    "    \n",
    "    times = u.math.arange(0*u.ms, num_steps*brainstate.environ.get_dt(), brainstate.environ.get_dt())\n",
    "    indices = u.math.arange(times.size)\n",
    "    predictions = brainstate.transform.for_loop(step, times, indices)\n",
    "    return compute_loss(predictions, targets)\n",
    "\n",
    "# Get trainable parameters\n",
    "params = network.states(brainstate.ParamState)\n",
    "\n",
    "# Compute gradients\n",
    "if len(params) > 0:\n",
    "    optimizer = braintools.optim.Adam(lr=1e-3)\n",
    "    grads, loss = brainstate.transform.grad(\n",
    "        loss_fn,\n",
    "        grad_states=params,\n",
    "        return_value=True\n",
    "    )()\n",
    "    print(f\"Loss: {loss}\")\n",
    "    # Update parameters with optimizer (if defined)\n",
    "    optimizer.update(grads)\n",
    "else:\n",
    "    # If no trainable parameters, just compute loss\n",
    "    loss = loss_fn()\n",
    "    print(f\"Loss (no trainable params): {loss}\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss (no trainable params): 0.0\n"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Physical Units System\n",
    "\n",
    "### Integration with brainunit\n",
    "\n",
    "``brainpy.state`` integrates ``brainunit`` for scientific accuracy:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.386693Z",
     "start_time": "2025-11-13T09:31:09.382194Z"
    }
   },
   "source": [
    "# Define with units\n",
    "tau = 10 * u.ms\n",
    "threshold = -50 * u.mV\n",
    "current = 5 * u.nA\n",
    "\n",
    "# Units are checked automatically\n",
    "neuron = brainpy.state.LIF(100, tau=tau, V_th=threshold)"
   ],
   "outputs": [],
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Benefits:\n",
    "\n",
    "- Prevents unit errors (e.g., ms vs s)\n",
    "- Self-documenting code\n",
    "- Automatic unit conversions\n",
    "- Scientific correctness"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unit Operations"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.432404Z",
     "start_time": "2025-11-13T09:31:09.427692Z"
    }
   },
   "source": [
    "# Arithmetic with units\n",
    "total_time = 100 * u.ms + 0.5 * u.second  # → 600 ms\n",
    "\n",
    "# Unit conversion\n",
    "time_in_seconds = (100 * u.ms).to_decimal(u.second)  # → 0.1\n",
    "\n",
    "# Unit checking (automatic in BrainPy operations)\n",
    "voltage = -65 * u.mV\n",
    "current = 2 * u.nA\n",
    "resistance = voltage / current  # Automatically gives MΩ"
   ],
   "outputs": [],
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ecosystem Integration\n",
    "\n",
    "``brainpy.state`` integrates tightly with its ecosystem:\n",
    "\n",
    "### braintools\n",
    "\n",
    "Utilities and tools:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.471630Z",
     "start_time": "2025-11-13T09:31:09.465919Z"
    }
   },
   "source": [
    "# Optimizers\n",
    "optimizer = braintools.optim.Adam(lr=1e-3)\n",
    "\n",
    "# Initializers\n",
    "init = braintools.init.KaimingNormal()\n",
    "\n",
    "# Surrogate gradients\n",
    "spike_fn = braintools.surrogate.ReluGrad()\n",
    "\n",
    "# Metrics (example with dummy data)\n",
    "# pred = jnp.array([0.1, 0.9])\n",
    "# target = jnp.array([0, 1])\n",
    "# loss = braintools.metric.cross_entropy(pred, target)"
   ],
   "outputs": [],
   "execution_count": 17
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### brainunit\n",
    "\n",
    "Physical units:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.496064Z",
     "start_time": "2025-11-13T09:31:09.490951Z"
    }
   },
   "source": [
    "# All standard SI units\n",
    "time = 10 * u.ms\n",
    "voltage = -65 * u.mV\n",
    "current = 2 * u.nA"
   ],
   "outputs": [],
   "execution_count": 18
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### brainstate\n",
    "\n",
    "Core framework (used automatically):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-13T09:31:09.518362Z",
     "start_time": "2025-11-13T09:31:09.512404Z"
    }
   },
   "source": [
    "import brainstate\n",
    "\n",
    "# Module system\n",
    "class Net(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        pass\n",
    "\n",
    "# Compilation\n",
    "@brainstate.transform.jit\n",
    "def fn():\n",
    "    return 0\n",
    "\n",
    "# Transformations\n",
    "# result = brainstate.transform.for_loop(...)"
   ],
   "outputs": [],
   "execution_count": 19
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ecosystem-py",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
